import numpy as np
import matplotlib.pyplot as plt
import os

def plot_solution(file_path, grid_size, figure_dir, name):
    # 读取 CSV 文件
    solution = np.loadtxt(file_path)
    
    # 网格的维度
    n = grid_size
    
    # 生成 x 轴的坐标，区间为 [0, 1]
    x = np.linspace(0, 1, n+1)
    
    # 创建图像
    plt.figure(figsize=(8, 6))
    plt.plot(x, solution, linestyle='-', color='b')
    plt.title(f'Solution for Grid Size {grid_size}')
    plt.xlabel('X-axis')
    plt.ylabel('Solution Value')
    plt.grid(True)
    
    # 保存图像
    os.makedirs(figure_dir, exist_ok=True)
    figure_path = os.path.join(figure_dir, f'{name}_{grid_size}.png')
    plt.savefig(figure_path)
    plt.close()
    print(f"Figure saved to {figure_path}")

def main():
    output_dir = "output"
    figure_dir = "figure"
    
    # 遍历输出目录中的所有 CSV 文件
    for filename in os.listdir(output_dir):
        if filename.endswith(".csv"):
            file_path = os.path.join(output_dir, filename)
            name = filename.split('_')[1]
            grid_size = int(filename.split('_')[2].split('.')[0])
            plot_solution(file_path, grid_size, figure_dir, name)

if __name__ == "__main__":
    main()